package com.bigdata.auth;

import org.apache.commons.lang3.StringUtils;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.security.UserGroupInformation;

import java.net.URI;
import java.util.Properties;

/**
 * Kerberos认证工具类
 *
 * @author xin.ding
 * @date 2021/1/28 16:41
 */
public class KerberosUtil {

    /**
     * 默认协议
     */
    private static final String DEFAULT_PROTOCOL = "SASL_PLAINTEXT";

    /**
     * 创建Kerberos HDFS文件系统环境
     *
     * @param url        hdfs连接地址
     * @param principal  keytab principal 名称
     * @param krb5Path   krb5.conf路径
     * @param keytabPath keytab文件路径
     */
    public static FileSystem buildHdfsKerberosEnv(String url, String principal, String krb5Path, String keytabPath) {
        FileSystem fs;
        try {
            System.setProperty("java.security.krb5.conf", krb5Path);
            Configuration configuration = new Configuration();
            configuration.set("hadoop.security.authentication", "Kerberos");
            configuration.set("fs.hdfs.impl", "org.apache.hadoop.hdfs.DistributedFileSystem");
            configuration.set("dfs.client.use.datanode.hostname", "true");

            UserGroupInformation.setConfiguration(configuration);
            UserGroupInformation.loginUserFromKeytab(principal, keytabPath);
            fs = FileSystem.get(URI.create(url), configuration);
        } catch (Exception e) {
            throw new AuthException("HDFS Kerberos认证失败！", e);
        }
        return fs;
    }

    /**
     * 构建HiveKerberos 认证环境
     *
     * @param krb5Path   krb5.conf 路径
     * @param principal  kerberos principal
     * @param keytabPath keytab文件路径
     */
    public static void buildHiveKerberosEnv(String krb5Path, String principal, String keytabPath) {
        System.setProperty("java.security.krb5.conf", krb5Path);
        Configuration conf = new Configuration();
        conf.set("hadoop.security.authentication", "Kerberos");
        UserGroupInformation.setConfiguration(conf);
        try {
            UserGroupInformation.loginUserFromKeytab(principal, keytabPath);
        } catch (Exception e) {
            throw new AuthException("Hive Kerberos 认证失败！", e);
        }
    }

    /**
     * 构建kafka Kerberos环境 采用默认协议,自动生成jaas文件
     *
     * @param properties kafka 配置信息
     * @param krb5Path   krb5文件路径
     * @param principal  kerberos principal 名称
     * @param keytab     kerberos keytab路径
     */
    public static void buildKafkaKerberosEnv(Properties properties, String krb5Path, String principal, String keytab) {
        buildKafkaKerberosEnv(properties, krb5Path, principal, keytab, null, DEFAULT_PROTOCOL);
    }

    /**
     * 构建kafka Kerberos环境
     *
     * @param properties kafka 配置信息
     * @param krb5Path   krb5文件路径
     * @param principal  kerberos principal 名称
     * @param keytab     kerberos keytab路径
     * @param jaasPath   jaas配置文件路径
     * @param protocol   认证协议
     */
    public static void buildKafkaKerberosEnv(Properties properties, String krb5Path, String principal, String keytab, String jaasPath, String protocol) {
        System.setProperty("sun.security.krb5.debug", "true");
        System.setProperty("java.security.krb5.conf", krb5Path);
        if (StringUtils.isBlank(jaasPath)) {
            jaasPath = JaasFactory.buildJaasFile(SourceType.KAFKA, principal, keytab);
        }
        System.setProperty("java.security.auth.login.config", jaasPath);
        if (StringUtils.isBlank(protocol)) {
            protocol = DEFAULT_PROTOCOL;
        }
        properties.put("security.protocol", protocol);
        properties.put("sasl.mechanism", "GSSAPI");
        properties.put("sasl.kerberos.service.name", "kafka");
    }

}
